import torchvision
from torch.utils.tensorboard import SummaryWriter

dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, transform=dataset_transform, download=True)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, transform=dataset_transform, download=True)

# train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True)
# test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True)

# print(test_dataset[0])
# print(test_dataset.classes)

# img, target = test_dataset[0]
# print(img)
# print(target)
# print(test_dataset.classes[target])

# img.show()

# print(img.shape)

writer = SummaryWriter('p10')
for i in range(10):
    img, target = test_dataset[i]
    writer.add_image('test_set', img, i)

writer.close()